import express from "express";
import { randomUUID } from "node:crypto";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import * as vscode from "vscode";
import z from "zod";
export function CreateMcpServer(port: number = 9981) {
  const app = express();
  app.use(express.json());

  // Map to store transports by session ID
  const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};

  // Handle POST requests for client-to-server communication
  app.post("/mcp", async (req, res) => {
    // Check for existing session ID
    const sessionId = req.headers["mcp-session-id"] as string | undefined;
    let transport: StreamableHTTPServerTransport;

    if (sessionId && transports[sessionId]) {
      // Reuse existing transport
      transport = transports[sessionId];
    } else if (!sessionId && isInitializeRequest(req.body)) {
      // New initialization request
      transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => randomUUID(),
        onsessioninitialized: (sessionId) => {
          // Store the transport by session ID
          transports[sessionId] = transport;
        },
        // DNS rebinding protection is disabled by default for backwards compatibility. If you are running this server
        // locally, make sure to set:
        // enableDnsRebindingProtection: true,
        // allowedHosts: ['127.0.0.1'],
      });

      // Clean up transport when closed
      transport.onclose = () => {
        if (transport.sessionId) {
          delete transports[transport.sessionId];
        }
      };
      const server = new McpServer({
        name: "code-puppeteer",
        version: "1.0.0",
      });

      // ... set up server resources, tools, and prompts ...
      server.registerTool(
        "hello",
        {
          title: "hello",
          description: "显示 VsCode 通知消息",
          inputSchema: {
            content: z.string().describe("通知内容"),
          },
        },
        ({ content }) => {
          vscode.window.showInformationMessage(content);
          return {
            content: [
              {
                type: "text",
                text: `通知内容: ${content}`,
              },
            ],
          };
        }
      );

      // 在编辑器当前光标处插入代码
      server.registerTool(
        "insert",
        {
          title: "insert",
          description: "在编辑器当前光标处插入代码",
          inputSchema: {
            content: z.string().describe("要插入的代码"),
            save: z.boolean().default(false).describe("是否保存文件"),
            format: z.boolean().default(false).describe("是否格式化代码"),
          },
        },
        async ({ content, save, format }) => {
          const editor = vscode.window.activeTextEditor;
          if (editor) {
            editor.edit((editBuilder) => {
              editBuilder.insert(editor.selection.active, content);
            });
            if (save) {
              await editor.document.save();
            }
            if (format) {
              const document = editor.document;
              try {
                // 获取文档的格式化编辑
                const textEdits = await vscode.commands.executeCommand<
                  vscode.TextEdit[]
                >("vscode.executeFormatDocumentProvider", document.uri);

                if (textEdits && textEdits.length > 0) {
                  // 应用格式化编辑
                  const edit = new vscode.WorkspaceEdit();
                  textEdits.forEach((textEdit) => {
                    edit.replace(
                      document.uri,
                      textEdit.range,
                      textEdit.newText
                    );
                  });

                  await vscode.workspace.applyEdit(edit);
                  vscode.window.showInformationMessage("文档格式化完成");
                }
              } catch (error) {
                vscode.window.showErrorMessage(`格式化失败: ${error}`);
              }
            }
          }
          return {
            content: [
              {
                type: "text",
                text: `已插入代码: ${content}`,
              },
            ],
          };
        }
      );

      // Connect to the MCP server
      await server.connect(transport);
    } else {
      // Invalid request
      res.status(400).json({
        jsonrpc: "2.0",
        error: {
          code: -32000,
          message: "Bad Request: No valid session ID provided",
        },
        id: null,
      });
      return;
    }

    // Handle the request
    await transport.handleRequest(req, res, req.body);
  });

  // Reusable handler for GET and DELETE requests
  const handleSessionRequest = async (
    req: express.Request,
    res: express.Response
  ) => {
    const sessionId = req.headers["mcp-session-id"] as string | undefined;
    if (!sessionId || !transports[sessionId]) {
      res.status(400).send("Invalid or missing session ID");
      return;
    }

    const transport = transports[sessionId];
    await transport.handleRequest(req, res);
  };

  // Handle GET requests for server-to-client notifications via SSE
  app.get("/mcp", handleSessionRequest);

  // Handle DELETE requests for session termination
  app.delete("/mcp", handleSessionRequest);

  return app.listen(port, () => {
    console.log(`MCP server is running on port ${port}`);
  });
}
